Fix dtype promotion issue in scalar indexing #32460
Conversation
jakevdp
force-pushed
the
fix-index-dtypes
branch
from
2487c5b to
9cbcc4b
Compare
|
/gemini review |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request fixes a dtype promotion issue in scalar indexing by ensuring that indices are promoted to a dtype suitable for the array being indexed, rather than just based on the index types themselves. The changes in jax/_src/numpy/indexing.py and jax/_src/checkify.py are robust and address the issue. A regression test is also added. I've found one potential issue where the out-of-bounds check might not respect the newly determined dtype, which could lead to incorrect error checking in some cases.
| index_dtype = lax_utils.int_dtype_for_shape(arr.shape, signed=True) | ||
| start_indices = [lax.convert_element_type(idx, index_dtype) for idx in start_indices] |
There was a problem hiding this comment.
This is a great change to ensure start_indices have a robustly determined dtype. However, the very next call to jnp_error._check_precondition_oob_dynamic_slice seems to counteract this by hardcoding int32 for its index checks.
This could lead to incorrect out-of-bounds checks if lax_utils.int_dtype_for_shape selected a larger type like int64 (e.g., for an array with large dimensions), as the indices would be downcast to int32 inside the check function, potentially causing overflow or incorrect comparisons.
To make this fix fully effective, _check_precondition_oob_dynamic_slice should probably be updated to respect the dtype of the incoming start_indices, similar to the changes made in checkify.py in this PR.
1 participant
Fixes #31396